package org.apache.ocean;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;

import org.apache.lucene.index.Term;
import org.apache.lucene.search.Query;

/**
 *
 * @author jasonr
 */
// TODO: unit test MultiShardDocFreq
public class MultiShardTermFrequency {
	public static Logger log = Logger.getLogger(MultiShardTermFrequency.class.getName());
  Query rewrittenQuery;
  ThreadPoolExecutor threadPool;
  List<SearchIndexShard> searchIndexShards;
  TimeoutState timeoutState;
  
  public MultiShardTermFrequency(Query rewrittenQuery, List<SearchIndexShard> searchIndexShards, ThreadPoolExecutor threadPool, TimeoutState timeoutState) {
    this.rewrittenQuery = rewrittenQuery;
    this.searchIndexShards = searchIndexShards;
    this.threadPool = threadPool;
    this.timeoutState = timeoutState;
  }
  
  public TermFrequencies getTermFrequencies() throws Exception {
    Set<Term> terms = new HashSet<Term>();
    rewrittenQuery.extractTerms(terms);
    TermFrequencies termFrequencies = new TermFrequencies(terms.size());
    for (Term term : terms) termFrequencies.put(term, new AtomicLong(0));
    List<Callable<Object>> docFreqTasks = new ArrayList<Callable<Object>>(searchIndexShards.size());
    for (SearchIndexShard searchIndexShard : searchIndexShards) {
      List<MultiIndexSearcher> indexSearchers = searchIndexShard.createMultiIndexSearchers();
      docFreqTasks.add(new TermFrequencyTask(indexSearchers, termFrequencies, timeoutState));
    }
    List<Future<Object>> futures = threadPool.invokeAll(docFreqTasks);
    for (Future future : futures) {
      try {
        if (!timeoutState.isOK()) throw new GeneralTimeoutException(timeoutState);
        future.get();
      } catch (ExecutionException executionException) {
        //throw (Exception)executionException.getCause();
        ClusterUtil.throwException(executionException);
      }
    }
    return termFrequencies;
  }
  
  public static class TermFrequencyTask implements Callable {
    List<MultiIndexSearcher> indexSearchers;
    TermFrequencies termFrequencies;
    TimeoutState timeoutState;
    
    public TermFrequencyTask(List<MultiIndexSearcher> indexSearchers, TermFrequencies termFrequencies, TimeoutState timeoutState) {
      this.indexSearchers = indexSearchers;
      this.termFrequencies = termFrequencies;
      this.timeoutState = timeoutState;
    }
    
    public Object call() throws Exception {
      for (MultiIndexSearcher indexSearcher : indexSearchers) {
        for (Map.Entry<Term,AtomicLong> entry : termFrequencies.entrySet()) {
          Term term = entry.getKey();
          AtomicLong frequency = entry.getValue();
          frequency.addAndGet(indexSearcher.docFreq(term));
          if (!timeoutState.isOK()) throw new GeneralTimeoutException(timeoutState);
        }
      }
      return null;
    }
  }
}
